import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns

# 假设 agent_positions 是一个 (A, T, 2) 的数组
# A: 车辆数量, T: 时间步数, 2: 车辆的 (x, y) 坐标
A = 100  # 车辆数量
T = 50   # 时间步数
np.random.seed(0)

# 随机生成车辆的位置数据
agent_positions = np.random.rand(A, T, 2) * 100  # 车辆的位置在 0 到 100 的区域内

# 选择多个时间步 (例如 t=10 到 t=30)，合并这些时间步的数据
start_time = 10
end_time = 30
positions_at_multiple_times = agent_positions[:, start_time:end_time, :].reshape(-1, 2)

# 提取所有时间步中车辆的 (x, y) 坐标
x_multiple = positions_at_multiple_times[:, 0]
y_multiple = positions_at_multiple_times[:, 1]

# 定义网格大小，这里设置为 10x10 的网格
grid_size = 10  # 可以根据需要调整

# 计算热力图
heatmap_multi, xedges_multi, yedges_multi = np.histogram2d(x_multiple, y_multiple, bins=grid_size, range=[[0, 100], [0, 100]])

# 绘制热力图
plt.figure(figsize=(8, 6))
sns.heatmap(heatmap_multi.T, cmap='YlGnBu', cbar=True, xticklabels=xedges_multi.astype(int), yticklabels=yedges_multi.astype(int))

# 添加标题和标签
plt.title(f'Heatmap of Vehicle Distribution from Time {start_time} to {end_time}')
plt.xlabel('X Position')
plt.ylabel('Y Position')

# 显示图形
plt.show()
